import argparse
import os
import platform
import re
import sqlalchemy
import subprocess
import zipfile

from sqlalchemy import Engine

import constants

from collections import deque, defaultdict
from openpyxl.reader.excel import load_workbook
from openpyxl.workbook import Workbook
from openpyxl.styles import Font
from pathlib import Path
from typing import List, Dict, Deque
from data_model import LaunchTraceData, LaunchData, ThreadTraceListData, ThreadTraceTreeData
from query_data import get_trace_frame_slice_duration, query_thread_trace_tree_list, query_launch_trace_list


def convert_time_to_ms(time) -> float:
    return time / constants.NS_TO_MS

def get_trace_streamer(system: str):
    if system == 'Windows':
        return os.path.join('trace_streamer_binary', 'trace_streamer_windows.exe')
    elif system == 'Linux':
        return os.path.join('trace_streamer_binary', 'trace_streamer_linux')
    elif system == 'Darwin':
        return os.path.join('trace_streamer_binary', 'trace_streamer_mac')
    else:
        raise Exception('Unknown OS System')


def parse_launch_trace_to_db(trace_path: str):
    if trace_path.endswith('.db'):
        return trace_path
    if trace_path.endswith('.insight') or trace_path.endswith('.ftrace') or trace_path.endswith('.htrace'):
        if not os.path.exists(constants.WORKSPACE):
            os.makedirs(constants.WORKSPACE)
        file_path = trace_path
        if trace_path.endswith('.insight'):
            with zipfile.ZipFile(trace_path, 'r') as zip_ref:
                htrace_file = [f for f in zip_ref.namelist() if f.endswith('.htrace')][0]
                zip_ref.extract(htrace_file, constants.WORKSPACE)
                file_path = os.path.join(constants.WORKSPACE, htrace_file)

        print('Parsing Launch Trace File: ' + file_path)
        filename = Path(file_path).stem
        db_path = "{}/{}".format(constants.WORKSPACE, filename + '.db')

        trace_streamer = get_trace_streamer(platform.system())
        
        # os.system('{} {} -e {}'.format(trace_streamer, file_path, db_path))
        try:
            result = subprocess.run([trace_streamer, file_path, '-e', db_path], check=True, text=True, capture_output=True)
        except subprocess.CalledProcessError as e:
            if not os.path.exists(file_path):
                print('File Not Found: ' + file_path)
            else:
                print('parse error: ' + str(e))
            exit(1)

        # print('Pass Success!')
        if trace_path.endswith('.insight'):
            os.remove(file_path)
        return db_path

    else:
        raise Exception('Not Supported File Type: {}'.format(trace_path))


def get_process_id(launch_trace_list: List[LaunchTraceData], bundle_name: str):
    for launch_trace in launch_trace_list:
        if bundle_name.endswith(launch_trace.thread_name):
            return launch_trace.thread_id
    return None


def create_launch_lifecycle_phase(launch_trace: LaunchTraceData, lifecycle_name: str):
    return LaunchData(lifecycle_name=lifecycle_name, start_time=launch_trace.start_time,
                      lifecycle_duration=launch_trace.duration, trace_name=launch_trace.name,
                      trace_duration=launch_trace.duration, thread_id=launch_trace.thread_id, children=None)


def set_process_create(launch_trace: LaunchTraceData, launch_data_array: List[LaunchTraceData],
                       thread_id_array: List[int]):
    launch_data_array[0] = create_launch_lifecycle_phase(launch_trace, constants.LIFE_CYCLE_NAME_PROCESS_CREATING)
    thread_id_array[0] = launch_trace.thread_id


def deal_process_create(launch_trace: LaunchTraceData, launch_data_array: List[LaunchTraceData],
                        thread_id_array: List[int]):
    is_same_event = (launch_data_array[0] is not None
                     and launch_trace.thread_id == launch_data_array[0].thread_id
                     and launch_trace.start_time < launch_data_array[0].start_time + launch_data_array[0].duration)
    if is_same_event:
        return
    else:
        set_process_create(launch_trace, launch_data_array, thread_id_array)


def parse_process_create(launch_trace: LaunchTraceData, bundle_name: str, launch_data_array: List[LaunchTraceData],
                         thread_id_array: List[int]):
    if constants.TRACE_START_ABILITY in launch_trace.name:
        set_process_create(launch_trace, launch_data_array, thread_id_array)
        return False
    if constants.TRACE_START_ABILITY_LOCKED in launch_trace.name and bundle_name.endswith(launch_trace.name):
        deal_process_create(launch_trace, launch_data_array, thread_id_array)
        return True
    if constants.TRACE_LOAD_ABILITY in launch_trace.name:
        if thread_id_array[0] == 0:
            deal_process_create(launch_trace, launch_data_array, thread_id_array)
        return True
    return False


def parse_application_launching(launch_trace: LaunchTraceData, bundle_name: str,
                                launch_data_array: List[LaunchTraceData],
                                thread_id_array: List[int]):
    if constants.TRACE_ATTACH_APPLICATION in launch_trace.name and bundle_name.endswith(launch_trace.name):
        launch_data_array[1] = create_launch_lifecycle_phase(launch_trace,
                                                             constants.LIFE_CYCLE_NAME_APPLICATION_LAUNCHING)
        thread_id_array[1] = launch_trace.thread_id
        return True
    return False


def parse_ui_ability_launching(launch_trace: LaunchTraceData, launch_data_array: List[LaunchTraceData],
                               process_id: int):
    if process_id == launch_trace.thread_id and constants.TRACE_LAUNCH_ABILITY in launch_trace.name:
        launch_data_array[2] = create_launch_lifecycle_phase(launch_trace, constants.LIFE_CYCLE_NAME_UI_ABILITY_LAUNCH)
        return True
    return False


def parse_ui_ability_on_foreground(launch_trace: LaunchTraceData, launch_data_array: List[LaunchTraceData],
                                   process_id: int, ability_name: str):
    if process_id == launch_trace.thread_id and constants.TRACE_FOREGROUND_ABILITY in launch_trace.name and ability_name in launch_trace.name:
        launch_data_array[3] = create_launch_lifecycle_phase(launch_trace,
                                                             constants.LIFE_CYCLE_NAME_UI_ABILITY_ON_FOREGROUND)
        return True
    return False


def get_frame_relate_key(trace_name: str):
    matcher = constants.PATTERN_FRAME_TRACE_RELATE.search(trace_name)
    if matcher:
        index_data = matcher.group().replace(" ", "")
        index_data_split = index_data.split(",")
        if len(index_data_split) == 3:
            return index_data_split[0] + ',' + index_data_split[2]
        return index_data
    return None


def parse_app_first_frame_rendering(launch_trace: LaunchTraceData, launch_data_array: List[LaunchTraceData],
                                    process_id: int):
    if not process_id == launch_trace.thread_id:
        return False

    if constants.TRACE_RECEIVE_VSYNC in launch_trace.name:
        launch_data_array[4] = create_launch_lifecycle_phase(launch_trace,
                                                             constants.LIFE_CYCLE_NAME_FIRST_APP_FRAME_RENDERING)
        return False
    if launch_data_array[
        4] is not None and constants.TRACE_APP_RENDERING in launch_trace.name and launch_trace.start_time < \
            launch_data_array[4].start_time + launch_data_array[4].trace_duration:
        relate_key = get_frame_relate_key(launch_trace.name)
        if relate_key is None:
            return False
        launch_data_array[4].name = relate_key
        return True
    return False


def parse_rs_first_frame_rendering(launch_trace: LaunchTraceData, launch_data_array: List[LaunchData],
                                   db_engine: Engine):
    if not constants.PROCESS_RENDER_SERVICE == launch_trace.thread_name:
        return False
    if constants.TRACE_RECEIVE_VSYNC in launch_trace.name:
        launch_data_array[5] = create_launch_lifecycle_phase(launch_trace,
                                                             constants.LIFE_CYCLE_NAME_FIRST_RS_FRAME_RENDERING)
        return False
    if launch_data_array[5] is not None and launch_data_array[4].trace_name in launch_trace.name and constants.TRACE_RS_RENDERING in launch_trace.name:
        duration = get_trace_frame_slice_duration(db_engine, launch_data_array[5].trace_name,
                                                  launch_data_array[5].thread_id, launch_data_array[5].start_time)
        if duration is not None and -1 != duration:
            launch_data_array[5].lifecycle_duration = duration
        return True
    return False

def parse_flush_trace_time_in_start_ability(launch_trace: LaunchTraceData, launch_data_array: List[LaunchTraceData], process_id: int, ability_name: str, db_engine: Engine):
    if process_id == launch_trace.thread_id and constants.TRACE_RECEIVE_VSYNC in launch_trace.name:
        tree_list = query_thread_trace_tree_list(db_engine, process_id, process_id, launch_trace.start_time, launch_trace.start_time + launch_trace.duration)
        for trace in tree_list:
            if constants.COLD_START_FINISH_SIGN in trace.name:
                launch_data_array[6] = create_launch_lifecycle_phase(launch_trace, ability_name)
                return True
    return False

def calculate_launch_lifecycle_duration(has_transit_array: List[bool], launch_data_array: List[LaunchData]):
    for i in range(3, -1, -1):
        if launch_data_array[i] is None:
            continue
        if launch_data_array[i].lifecycle_duration < -1:
            launch_data_array[i] = None
            continue
        if not has_transit_array[i] and not launch_data_array[i].lifecycle_duration == -1:
            continue
        for j in range(i + 1, 5):
            if launch_data_array[j] is None:
                continue
            launch_data_array[i].lifecycle_duration = launch_data_array[j].start_time - launch_data_array[i].start_time
            if launch_data_array[i].lifecycle_duration < 0:
                launch_data_array[i] = None
            break

    # EntryAbility二刷lifecycle
    if len(launch_data_array) == 7 and launch_data_array[6] and launch_data_array[5]:
        prev_end_time = launch_data_array[5].start_time + launch_data_array[5].lifecycle_duration
        launch_data_array[6].lifecycle_duration = launch_data_array[6].start_time - prev_end_time + launch_data_array[6].trace_duration
        launch_data_array[6].start_time = prev_end_time


def parse_launch_life_cycle_phase(launch_trace_list: List[LaunchTraceData], bundle_name: str, ability_name: str,
                                  db_engine: Engine):
    if constants.COLD_START_FINISH_SIGN is not None:
        life_cycle_phase_num = 7
    else:
        life_cycle_phase_num = 6
    has_transit_array = [False] * life_cycle_phase_num
    launch_data_array = [None] * life_cycle_phase_num
    thread_id_array = [0] * 4
    process_id = get_process_id(launch_trace_list, bundle_name)
    thread_id_array[2] = process_id

    for launch_trace in launch_trace_list:
        # 解析process create
        if not has_transit_array[0]:
            has_transit_array[0] = parse_process_create(launch_trace, bundle_name, launch_data_array, thread_id_array)
        # 解析application launch
        if not has_transit_array[1]:
            has_transit_array[1] = parse_application_launching(launch_trace, bundle_name, launch_data_array,
                                                               thread_id_array)
        # 解析ui application launch
        if not has_transit_array[2]:
            has_transit_array[2] = parse_ui_ability_launching(launch_trace, launch_data_array, process_id)
        # 解析ui onforeground
        if not has_transit_array[3]:
            has_transit_array[3] = parse_ui_ability_on_foreground(launch_trace, launch_data_array, process_id,
                                                                  ability_name)
        # 解析app frame
        if not has_transit_array[4]:
            has_transit_array[4] = parse_app_first_frame_rendering(launch_trace, launch_data_array, process_id)
        # 解析render phase
        if has_transit_array[4] and not has_transit_array[5]:
            has_transit_array[5] = parse_rs_first_frame_rendering(launch_trace, launch_data_array, db_engine)

        # 解析二刷场景
        if life_cycle_phase_num == 7 and has_transit_array[5] and not has_transit_array[6]:
            has_transit_array[6] = parse_flush_trace_time_in_start_ability(launch_trace, launch_data_array, process_id, ability_name, db_engine)

    calculate_launch_lifecycle_duration(has_transit_array, launch_data_array)
    return launch_data_array


def parse_thread_trace_tree_data(thread_trace_list_data: ThreadTraceListData, start_time: int,
                                 end_time: int) -> ThreadTraceTreeData:
    start = max(thread_trace_list_data.start_time, start_time)
    end = min(thread_trace_list_data.start_time + thread_trace_list_data.duration, end_time)
    return ThreadTraceTreeData(
        name=thread_trace_list_data.name,
        start_time=start,
        duration=end - start,
        depth=thread_trace_list_data.depth,
        children=[]
    )


def build_thread_trace_tree_recursive(base_node_list: Deque[ThreadTraceTreeData],
                                      thread_trace_map_by_depth: Dict[int, Deque], depth: int) -> List[
    ThreadTraceTreeData]:
    result = list(base_node_list)
    children_list = thread_trace_map_by_depth[depth]
    if not children_list or not base_node_list:
        return result
    base_node = base_node_list.popleft()
    for child in children_list:
        while child.start_time + child.duration < base_node.start_time or child.start_time > base_node.start_time + base_node.duration:
            if not base_node_list:
                base_node = None
                break
            base_node = base_node_list.popleft()
        if base_node is None:
            break
        if not base_node.children:
            base_node.children = [child]
        else:
            base_node.children.append(child)
    build_thread_trace_tree_recursive(children_list, thread_trace_map_by_depth, depth + 1)
    return result


def add_children_to_thread_trace_tree(thread_trace_tree_data_list: List[ThreadTraceTreeData]) -> ThreadTraceTreeData:
    thread_trace_map_by_depth = defaultdict(deque)
    for data in thread_trace_tree_data_list:
        thread_trace_map_by_depth[data.depth].append(data)
    root_list = thread_trace_map_by_depth[0]
    trace_tree_data = build_thread_trace_tree_recursive(root_list, thread_trace_map_by_depth, 1)
    if not trace_tree_data:
        return []
    else:
        return trace_tree_data


def build_thread_trace_tree(trace_tree_list: List[ThreadTraceListData], thread_id: int, start_time: int,
                            end_time: int) -> list:
    if not trace_tree_list:
        return []

    target_thread_trace_data_list = [
        parse_thread_trace_tree_data(data, start_time, end_time) for data in trace_tree_list if data.tid == thread_id
    ]

    target_thread_trace_data_list.sort(key=lambda x: x.start_time)

    thread_trace_tree = []
    thread_trace_tree.append(add_children_to_thread_trace_tree(target_thread_trace_data_list))
    return thread_trace_tree

# 排除一些data数字不同导致name不匹配 以及一些打点开关导致的不一致 如 tracename|M32
def compare_name_without_number(name1: str, name2: str) -> bool:
    replace_num_name_1 = re.sub("\d", "", name1).split('|')[0]
    replace_num_name_2 = re.sub("\d", "", name2).split('|')[0]
    return replace_num_name_1 == replace_num_name_2

def max_match_coverage(list1: List[ThreadTraceListData], list2: List[ThreadTraceListData]) -> List[List[ThreadTraceListData]]:
    m = len(list1)
    n = len(list2)
    dp = [[0 for _ in range(n + 1)] for _ in range(m + 1)]

    for i in range(1, m + 1):
        for j in range(1, n + 1):
            if compare_name_without_number(list1[i - 1].name, list2[j - 1].name):
                # 可能有些存在回调不一致导致时序有些问题匹配不上，以耗时为准来找到耗时最长的顺序匹配路径
                dp[i][j] = dp[i - 1][j - 1] + list1[i - 1].duration + list2[j - 1].duration
            else:
                dp[i][j] = max(dp[i - 1][j], dp[i][j - 1])

    i, j = m, n

    max_coverage = []
    while i > 0 and j > 0:
        if compare_name_without_number(list1[i - 1].name, list2[j - 1].name):
            max_coverage.append([list1[i - 1], list2[j - 1]])
            i -= 1
            j -= 1
        elif dp[i - 1][j] > dp[i][j - 1]:
            i -= 1
        else:
            j -= 1
    max_coverage.reverse()
    return max_coverage

def compare_trace_tree(launch_trace_tree_1: List[ThreadTraceTreeData], launch_trace_tree_2: List[ThreadTraceTreeData],
                       depth: int, ws, launch_lifecycle_phase_name):
    queue_1 = deque(launch_trace_tree_1)
    queue_2 = deque(launch_trace_tree_2)

    for dep in range(depth):
        prev_trace_endTime_1 = None
        prev_trace_endTime_2 = None
        if not queue_1 or not queue_2:
            break
        queue_1_len = len(queue_1)
        queue_2_len = len(queue_2)
        assert queue_1_len == queue_2_len, "The number of nodes at depth {} is not equal: {} vs {}".format(dep,
                                                                                                           queue_1_len,
                                                                                                           queue_2_len)
        for _ in range(queue_1_len):
            node_1 = queue_1.popleft()
            node_2 = queue_2.popleft()
            matched_pairs = max_match_coverage(node_1, node_2)
            for idx, (item1, item2) in enumerate(matched_pairs):
                # 目前每个depth最开始的prev diff 不考虑
                if prev_trace_endTime_1 and prev_trace_endTime_2:
                    ws.append([launch_lifecycle_phase_name, item1.name,
                                     convert_time_to_ms(item1.start_time),
                                     convert_time_to_ms(item1.duration),
                                     convert_time_to_ms(item2.start_time),
                                     convert_time_to_ms(item2.duration),
                                     convert_time_to_ms(item2.duration - item1.duration),
                                     convert_time_to_ms(item1.start_time - prev_trace_endTime_1),
                                     convert_time_to_ms(item2.start_time - prev_trace_endTime_2), dep + 1])
                else:
                    ws.append([launch_lifecycle_phase_name, item1.name,
                               convert_time_to_ms(item1.start_time),
                               convert_time_to_ms(item1.duration),
                               convert_time_to_ms(item2.start_time),
                               convert_time_to_ms(item2.duration),
                               convert_time_to_ms(item2.duration - item1.duration),
                               "-",
                               "-", dep + 1])
                prev_trace_endTime_1 = item1.start_time + item1.duration
                prev_trace_endTime_2 = item2.start_time + item2.duration
            for item1, item2 in matched_pairs:
                if item1.children or item2.children:
                    queue_1.append(item1.children)
                    queue_2.append(item2.children)

def is_number(value):
    try:
        float(value)
        return True
    except (ValueError, TypeError):
        return False

def highlight_threshold_time(output_path: str):
    wb = load_workbook(output_path)
    ws = wb['lifecycle_compare_detail']
    for row in ws.iter_rows():
        if is_number(row[6].value):
            if float(row[6].value) >= constants.TIME_THRESHOLD:
                row[6].font = Font(color='FF0000')
            elif float(row[6].value) <= -constants.TIME_THRESHOLD:
                row[6].font = Font(color='00FF00')
        if is_number(row[7].value):
            if float(row[7].value) > constants.TIME_THRESHOLD:
                row[7].font = Font(color='FF0000')
            elif float(row[7].value) <= -constants.TIME_THRESHOLD:
                row[7].font = Font(color='00FF00')
        if is_number(row[8].value):
            if float(row[8].value) > constants.TIME_THRESHOLD:
                row[8].font = Font(color='FF0000')
            elif float(row[8].value) <= -constants.TIME_THRESHOLD:
                row[8].font = Font(color='00FF00')
    wb.save(output_path)


def compare_trace_file(trace_file_1: str, trace_file_2: str, bundle_name: str, ability_name: str, depth: int):
    db_path_1 = "sqlite:///{}".format(parse_launch_trace_to_db(trace_file_1))
    db_path_2 = "sqlite:///{}".format(parse_launch_trace_to_db(trace_file_2))

    engine_1 = sqlalchemy.create_engine(db_path_1)
    engine_2 = sqlalchemy.create_engine(db_path_2)

    launch_trace_list_1 = query_launch_trace_list(engine_1)
    launch_trace_list_2 = query_launch_trace_list(engine_2)

    # 根据打点数据获得冷启动阶段
    launch_lifecycle_1 = parse_launch_life_cycle_phase(launch_trace_list_1, bundle_name, ability_name, engine_1)
    launch_lifecycle_2 = parse_launch_life_cycle_phase(launch_trace_list_2, bundle_name, ability_name, engine_2)

    process_id_1 = get_process_id(launch_trace_list_1, bundle_name)
    process_id_2 = get_process_id(launch_trace_list_2, bundle_name)

    if not os.path.exists(constants.OUTPUT_SPACE):
        os.makedirs(constants.OUTPUT_SPACE)
    output_file = os.path.join(constants.OUTPUT_SPACE, 'compare_trace_tree_depth_{}.xlsx'.format(depth))
    wb = Workbook()
    ws = wb.active
    ws.sheet_format.defaultColWidth = constants.DEFAULT_COL_WIDTH
    ws.title = 'lifecycle_compare'
    header = ["launch lifecycle phase", "trace1 startTime", "trace1 duration", "trace2 startTime",
              "trace2 duration", "duration diff"]
    ws.append(header)
    # 生成对应的冷启动阶段对比
    for i in range(len(launch_lifecycle_1)):
        if launch_lifecycle_1[i] is None or launch_lifecycle_2[i] is None:
            continue
        launch_lifecycle_phase_name = launch_lifecycle_1[i].lifecycle_name
        ws.append([launch_lifecycle_phase_name,
                         convert_time_to_ms(launch_lifecycle_1[i].start_time),
                         convert_time_to_ms(launch_lifecycle_1[i].lifecycle_duration),
                         convert_time_to_ms(launch_lifecycle_2[i].start_time),
                         convert_time_to_ms(launch_lifecycle_2[i].lifecycle_duration),
                         convert_time_to_ms(
                             launch_lifecycle_2[i].lifecycle_duration - launch_lifecycle_1[i].lifecycle_duration)])

    ws = wb.create_sheet('lifecycle_compare_detail')
    ws.sheet_format.defaultColWidth = constants.DEFAULT_COL_WIDTH

    # 继续对比每个阶段的线程打点树
    ws.append(["launch lifecycle phase", "trace name", "trace1 startTime", "trace1 duration",
                     "trace2 startTime", "trace2 duration", "duration diff", "trace1 prev diff",
                     "trace2 prev diff", "depth"])
    for i in range(len(launch_lifecycle_1)):
        if launch_lifecycle_1[i] is None or launch_lifecycle_2[i] is None:
            continue
        launch_lifecycle_phase_name = launch_lifecycle_1[i].lifecycle_name

        launch_trace_tree_list_1 = query_thread_trace_tree_list(engine_1, process_id_1, process_id_1,
                                                                launch_lifecycle_1[i].start_time,
                                                                launch_lifecycle_1[i].start_time +
                                                                launch_lifecycle_1[i].lifecycle_duration)
        launch_trace_tree_list_2 = query_thread_trace_tree_list(engine_2, process_id_2, process_id_2,
                                                                launch_lifecycle_2[i].start_time,
                                                                launch_lifecycle_2[i].start_time +
                                                                launch_lifecycle_2[i].lifecycle_duration)

        launch_trace_tree_1 = build_thread_trace_tree(launch_trace_tree_list_1, process_id_1,
                                                      launch_lifecycle_1[i].start_time,
                                                      launch_lifecycle_1[i].start_time + launch_lifecycle_1[
                                                          i].lifecycle_duration)
        launch_trace_tree_2 = build_thread_trace_tree(launch_trace_tree_list_2, process_id_2,
                                                      launch_lifecycle_2[i].start_time,
                                                      launch_lifecycle_2[i].start_time + launch_lifecycle_2[
                                                          i].lifecycle_duration)

        compare_trace_tree(launch_trace_tree_1, launch_trace_tree_2, depth, ws, launch_lifecycle_phase_name)
    wb.save(output_file)
    highlight_threshold_time(output_file)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('trace_file_1', help='输入文件绝对路径')
    parser.add_argument('trace_file_2', help='输入文件绝对路径')
    parser.add_argument('bundle_name', help='应用包名')
    parser.add_argument('depth', type=int, help='线程打点树深度')
    parser.add_argument('--name', help='应用入口Ability名称', default='EntryAbility')

    args = parser.parse_args()

    # compare_trace_file('Launch_ALN-AL80_com.alipay.mobile.client_MainProcess_20250629T162431.insight', 'Launch_ALN-AL80_com.alipay.mobile.client_MainProcess_20250630T205225.insight',
    #    'y.mobile.client', 'EntryAbility', 6)

    compare_trace_file(args.trace_file_1, args.trace_file_2, args.bundle_name, args.name, args.depth)
    print('Trace comparison completed. Output file: {}'.format(
        os.path.join(constants.OUTPUT_SPACE, 'compare_trace_tree_depth_{}.xlsx'.format(args.depth))))
